import torch
import torch.nn as nn
import pickle
import re
import os

# 导入原始文件中的 Word2Sequence 和 LSTMClassifier 类
# 确保 predict.py 和您的训练脚本在同一个目录下，或者调整导入路径
from model import Word2Sequence, LSTMClassifier, tokenlize


# ======================
#   加载词汇表
# ======================
def load_vocab(vocab_path="./model/ws.pkl"):
    with open(vocab_path, "rb") as f:
        ws = pickle.load(f)
    print(f"词汇表加载成功，大小: {len(ws)}")
    return ws


# ======================
#   加载模型
# ======================
def load_model(vocab_size, model_path="./model/lstm/best_model.pth", device="cuda"):
    model = LSTMClassifier(vocab_size=vocab_size)
    # 确保加载模型到正确的设备
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()  # 设置为评估模式
    print(f"模型加载成功: {model_path}")
    return model


# ======================
#   预测函数
# ======================
def predict_sentiment(text, model, ws, max_len=200, device="cuda"):
    # 文本预处理
    tokenized_text = tokenlize(text)
    # 将文本转换为数字序列
    sequence = ws.transform(tokenized_text, max_len=max_len)
    # 转换为 PyTorch tensor
    input_tensor = torch.tensor([sequence]).to(device)  # 增加一个batch维度

    with torch.no_grad():
        outputs = model(input_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        # 获取预测类别（0代表负面，1代表正面）
        predicted_class = torch.argmax(probabilities, dim=1).item()

    sentiment = "正面 (Positive)" if predicted_class == 1 else "负面 (Negative)"
    return sentiment, probabilities.cpu().numpy()[0]  # 返回概率，方便查看


if __name__ == '__main__':
    # 确定设备
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # 1. 加载词汇表
    ws = load_vocab()

    # 2. 实例化模型并加载权重
    # 这里的 vocab_size 必须和训练时使用的相同
    model = load_model(vocab_size=len(ws), device=device)

    # 3. 进行预测
    test_sentences = [
        "This movie was absolutely fantastic! I loved every minute of it.",
        "What a terrible film, utterly boring and a complete waste of time.",
        "It was okay, nothing special but not bad either.",
        "The acting was superb and the plot was engaging."
    ]

    print("\n--- 预测结果 ---")
    for sentence in test_sentences:
        sentiment, probabilities = predict_sentiment(sentence, model, ws, device=device)
        print(f"文本: '{sentence}'")
        print(f"预测情感: {sentiment}")
        print(f"概率: 负面={probabilities[0]:.4f}, 正面={probabilities[1]:.4f}")
        print("-" * 30)

    # 您还可以尝试输入自己的文本
    while True:
        user_input = input("\n请输入一段英文评论 (输入 'quit' 退出): ")
        if user_input.lower() == 'quit':
            break

        sentiment, probabilities = predict_sentiment(user_input, model, ws, device=device)
        print(f"预测情感: {sentiment}")
        print(f"概率: 负面={probabilities[0]:.4f}, 正面={probabilities[1]:.4f}")